package aipserver

import (
	"apedge/dao"
	"apedge/dto"
	"apedge/pkg/common"
	"apedge/pkg/configs"
	"apedge/pkg/database"
	"apedge/pkg/protocol"
	"apedge/pkg/utils"
	"apedge/remote"
	"archive/tar"
	"compress/gzip"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"math/big"
	"net"
	"os"
	"strconv"
	"time"

	"github.com/jinzhu/gorm"
)

// apedge

func GenNodeUuid() string {
	timeN := time.Now().UnixNano()
	return fmt.Sprintf("node-%s", strconv.FormatInt(timeN, 10))
}

func CheckNodeNameExist(tx *gorm.DB, name string) (bool, error) {
	var node dao.ApedgeNode
	var err error

	node.NodeName = name
	count, err := node.GetNodeNumByName(tx)
	if err != nil {
		return false, err
	}
	if count != 0 {
		return true, nil
	}
	return false, nil
}

func CheckApedgeInferenceNameExist(tx *gorm.DB, name string) (bool, error) {
	var node dao.ApedgeInference
	var err error

	node.Name = name
	count, err := node.GetNumByName(tx)
	if err != nil {
		return false, err
	}
	if count != 0 {
		return true, nil
	}
	return false, nil
}

func CheckNodeGroupNameExist(tx *gorm.DB, name string) (bool, error) {
	var node dao.ApedgeNodeGroup
	var err error

	node.Name = name
	count, err := node.GetNodeGroupNumByName(tx)
	if err != nil {
		return false, err
	}
	if count != 0 {
		return true, nil
	}
	return false, nil
}

func CheckNodeIdExist(tx *gorm.DB, nodeId string) (bool, error) {
	t := dao.ApedgeNode{
		NodeId: nodeId,
	}
	count, err := t.GetNodeNumByNodeId(tx)
	if err != nil {
		return false, err
	}

	if count != 0 {
		return true, nil
	}
	return false, nil
}

// CreateEdgeNodeCert create node certification
func CreateEdgeNodeCert(userInfo *protocol.UserInfoAAA, nodeId string) error {
	rsaDirPath := GetRsaDirectory(nodeId)
	privateKeyPath := GetPrivateKeyPath(nodeId)
	publicKeyPath := GetPublicKeyPath(nodeId)
	packPath := GetPackagePath(nodeId)

	fmt.Println("rsaDirPath: ", rsaDirPath)
	fmt.Println("privateKeyPath: ", privateKeyPath)
	fmt.Println("publicKeyPath: ", publicKeyPath)
	fmt.Println("packPath: ", packPath)

	// create dir
	err := os.MkdirAll(rsaDirPath, 0755)
	if err != nil {
		return err
	}
	// get ca config
	caCrtPath := configs.Config.Cluster.CACrtPath
	caKeyPath := configs.Config.Cluster.CAKeyPath
	certDuration := configs.Config.Cluster.CertDuration
	if certDuration <= 0 {
		certDuration = 365
	}
	caEncryptPass := configs.Config.Cluster.CAEncryptPass

	// create rsa key pair
	privateKey, _ := rsa.GenerateKey(rand.Reader, 2048)
	// logger.Infoln("gen privateKey: ", privateKey)
	// save private key file
	derStream := x509.MarshalPKCS1PrivateKey(privateKey)
	block := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: derStream,
	}
	privateKeyFile, err := os.Create(privateKeyPath)
	if err != nil {
		return err
	}
	defer privateKeyFile.Close()

	err = pem.Encode(privateKeyFile, block)
	if err != nil {
		return err
	}
	// create public key file
	publicKey := &privateKey.PublicKey
	derPkix, err := x509.MarshalPKIXPublicKey(publicKey)
	if err != nil {
		return err
	}
	block = &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: derPkix,
	}
	publicKeyFile, err := os.Create(publicKeyPath)
	if err != nil {
		return err
	}
	defer publicKeyFile.Close()

	err = pem.Encode(publicKeyFile, block)
	if err != nil {
		return err
	}
	// sign public key
	// read caCrt
	caFile, err := ioutil.ReadFile(caCrtPath)
	if err != nil {
		return err
	}

	caBlock, _ := pem.Decode(caFile)
	caCert, err := x509.ParseCertificate(caBlock.Bytes)
	if err != nil {
		return err
	}
	// copy crt file
	source, err := os.Open(caCrtPath)
	if err != nil {
		return err
	}
	defer source.Close()

	destination, err := os.Create(GetCACrtPath(nodeId))
	if err != nil {
		return err
	}

	defer destination.Close()
	_, err = io.Copy(destination, source)
	if err != nil {
		return err
	}
	// read ca key
	keyFile, err := ioutil.ReadFile(caKeyPath)
	if err != nil {
		return err
	}
	keyBlock, _ := pem.Decode(keyFile)
	decryBlock, err := x509.DecryptPEMBlock(keyBlock, []byte(caEncryptPass))
	if err != nil {
		return err
	}
	caKey, err := x509.ParsePKCS1PrivateKey(decryBlock)
	if err != nil {
		return err
	}

	// create request
	request := &x509.Certificate{
		SerialNumber: big.NewInt(1653),
		Subject: pkix.Name{
			Country:            []string{"CN"},
			Organization:       []string{"apulis"},
			OrganizationalUnit: []string{"apulisedge"},
			Locality:           []string{"ShenZhen"},
			Province:           []string{"ShenZhen"},
			CommonName:         "node-" + nodeId,
			Names:              []pkix.AttributeTypeAndValue{},
			ExtraNames:         []pkix.AttributeTypeAndValue{},
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().AddDate(0, certDuration, 0), //证书有效期结束时间
		BasicConstraintsValid: true,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageDataEncipherment,
	}
	//sign
	ca, err := x509.CreateCertificate(rand.Reader, request, caCert, publicKey, caKey)
	if err != nil {
		return err
	}

	// save cert
	certPath := GetCrtPath(nodeId)
	certFile, err := os.Create(certPath)
	if err != nil {
		return err
	}
	defer certFile.Close()

	caPem := &pem.Block{
		Type:  "CERTIFICATE",
		Bytes: ca,
	}
	err = pem.Encode(certFile, caPem)
	if err != nil {
		return err
	}

	// pack all file
	tarFile, err := os.Create(packPath)
	if err != nil {
		return err
	}
	defer tarFile.Close()
	// create writer
	gzipWriter := gzip.NewWriter(tarFile)
	defer gzipWriter.Close()
	tarWriter := tar.NewWriter(gzipWriter)
	defer tarWriter.Close()
	err = packFile(privateKeyPath, tarWriter)
	if err != nil {
		return err
	}
	err = packFile(certPath, tarWriter)
	if err != nil {
		return err
	}
	err = packFile(GetCACrtPath(nodeId), tarWriter)
	if err != nil {
		return err
	}

	return nil
}

func GetRsaDirectory(nodeId string) string {
	rsaDirPath := "/apedge/rsa/" + nodeId
	return rsaDirPath
}

func GetPrivateKeyPath(nodeUniqueName string) string {
	privateKeyPath := GetRsaDirectory(nodeUniqueName) + "/server.key"

	return privateKeyPath
}

func GetPublicKeyPath(nodeUniqueName string) string {
	publicKeyPath := GetRsaDirectory(nodeUniqueName) + "/server.pub"

	return publicKeyPath
}

func GetPackagePath(nodeUniqueName string) string {
	packPath := GetRsaDirectory(nodeUniqueName) + "/pack.tar.gz"

	return packPath
}
func GetCrtPath(nodeUniqueName string) string {
	packPath := GetRsaDirectory(nodeUniqueName) + "/server.crt"

	return packPath
}

func GetCACrtPath(nodeUniqueName string) string {
	packPath := GetRsaDirectory(nodeUniqueName) + "/rootCA.crt"

	return packPath
}

func packFile(srcFilePath string, tarWriter *tar.Writer) error {
	srcFile, err := os.Open(srcFilePath)
	if err != nil {
		return err
	}
	defer srcFile.Close()

	fileInfo, err := srcFile.Stat()
	if err != nil {
		return err
	}
	header, err := tar.FileInfoHeader(fileInfo, "")
	if err != nil {
		return err
	}
	err = tarWriter.WriteHeader(header)
	if err != nil {
		return err
	}
	_, err = io.Copy(tarWriter, srcFile)
	if err != nil {
		return err
	}
	return nil
}

func GenOutApedgeNodeItem(item dao.ApedgeNode) (dto.OutApedgeNodeItem, error) {
	out := dto.OutApedgeNodeItem{
		ID:                      item.ID,
		OrgId:                   item.OrgId,
		GroupId:                 item.GroupId,
		UserId:                  item.UserId,
		UserName:                item.UserName,
		NodeName:                item.NodeName,
		NodeId:                  item.NodeId,
		Describe:                item.Describe,
		NodeType:                item.NodeType,
		Arch:                    item.Arch,
		Cpu:                     item.Cpu,
		Mem:                     item.Mem,
		Status:                  item.Status,
		ContainerRuntimeVersion: item.ContainerRuntimeVersion,
		DeviceSeries:            item.DeviceSeries,
		OsImage:                 item.OsImage,
		CreatedAt:               utils.TimeToTimestamp(item.CreatedAt),
		UpdatedAt:               utils.TimeToTimestamp(item.UpdatedAt),
	}

	return out, nil
}

func GenOutApedgeNodeGroupItem(item dao.ApedgeNodeGroup) (dto.OutApedgeNodeGroupItem, error) {
	var onlineRatio string
	var onlineNodeNum int
	var out dto.OutApedgeNodeGroupItem

	nodeIds := []string{}
	err := json.Unmarshal([]byte(item.NodeIds), &nodeIds)
	if err != nil {
		return out, err
	}

	tx := database.Db
	nodeTemp := &dao.ApedgeNode{}
	nodeList, err := nodeTemp.GetNodeListByNodeIds(tx, nodeIds)
	if err != nil {
		return out, err
	}

	for _, n := range nodeList {
		if n.Status == common.NODE_STATUS_ONLINE {
			onlineNodeNum += 1
		}
	}
	onlineRatio = fmt.Sprintf("%d/%d", onlineNodeNum, len(nodeList))

	out = dto.OutApedgeNodeGroupItem{
		ID:          item.ID,
		OrgId:       item.OrgId,
		GroupId:     item.GroupId,
		UserId:      item.UserId,
		UserName:    item.UserName,
		Name:        item.Name,
		Describe:    item.Describe,
		NodeType:    item.Type,
		OnlineRatio: onlineRatio,
		CreatedAt:   utils.TimeToTimestamp(item.CreatedAt),
		UpdatedAt:   utils.TimeToTimestamp(item.UpdatedAt),
	}

	return out, nil
}

func GenOutApedgeInferenceItem(item dao.ApedgeInference) (dto.OutApedgeInferenceItem, error) {
	var models []dao.ApflowModel
	var nodeIds []string
	var containerPorts dto.ContainerPorts
	var envs map[string]string
	var runSvcNum int
	var svcNum int
	var nodesStatus string

	err := json.Unmarshal([]byte(item.Models), &models)
	if err != nil {
		return dto.OutApedgeInferenceItem{}, err
	}
	err = json.Unmarshal([]byte(item.NodeIds), &nodeIds)
	if err != nil {
		return dto.OutApedgeInferenceItem{}, err
	}
	err = json.Unmarshal([]byte(item.ContainerPorts), &containerPorts)
	if err != nil {
		return dto.OutApedgeInferenceItem{}, err
	}
	err = json.Unmarshal([]byte(item.Envs), &envs)
	if err != nil {
		return dto.OutApedgeInferenceItem{}, err
	}

	tx := database.Db
	for _, nodeId := range nodeIds {
		ns := dao.ApedgeServiceNode{
			NodeId:    nodeId,
			ServiceId: item.ServiceId,
		}
		nsList, err := ns.GetListByNodeIdAndServiceId(tx)
		if err != nil {
			return dto.OutApedgeInferenceItem{}, err
		}
		for _, one := range nsList {
			svcNum += 1
			if one.Status == common.TASK_STATUS_RUNNING {
				runSvcNum += 1
			}
		}
	}
	nodesStatus = fmt.Sprintf("%d/%d", runSvcNum, svcNum)

	out := dto.OutApedgeInferenceItem{
		ID:                     item.ID,
		OrgId:                  item.OrgId,
		GroupId:                item.GroupId,
		UserId:                 item.UserId,
		UserName:               item.UserName,
		ServiceId:              item.ServiceId,
		Name:                   item.Name,
		Describe:               item.Describe,
		Caller:                 item.Caller,
		StartTime:              item.StartTime,
		Url:                    item.Url,
		Status:                 item.Status,
		Models:                 models,
		Envs:                   envs,
		NodeIds:                nodeIds,
		NodesStatus:            nodesStatus,
		ContainerImage:         item.ContainerImage,
		ContainerImageTag:      item.ContainerImageTag,
		ContainerRunCmd:        item.ContainerRunCmd,
		ContainerRunParams:     item.ContainerRunParams,
		ContainerPrivilegeMode: item.ContainerPrivilegeMode,
		ContainerResart:        item.ContainerResart,
		ContainerNetwork:       item.ContainerNetwork,
		ContainerPorts:         containerPorts,
		CreatedAt:              utils.TimeToTimestamp(item.CreatedAt),
		UpdatedAt:              utils.TimeToTimestamp(item.UpdatedAt),
	}

	return out, nil
}

func GenOutServiceNodeItem(item dao.ApedgeServiceNode, nodeName string) (dto.OutServiceNodeItem, error) {

	out := dto.OutServiceNodeItem{
		ID:            item.ID,
		OrgId:         item.OrgId,
		GroupId:       item.GroupId,
		UserId:        item.UserId,
		UserName:      item.UserName,
		ServiceNodeId: item.ServiceNodeId,
		NodeId:        item.NodeId,
		NodeName:      nodeName,
		ServiceId:     item.ServiceId,
		Status:        item.Status,
		CreatedAt:     utils.TimeToTimestamp(item.CreatedAt),
		UpdatedAt:     utils.TimeToTimestamp(item.UpdatedAt),
	}
	return out, nil
}

func CreateInstallScripts(userInfo *protocol.UserInfoAAA, NodeId string) ([]string, error) {
	var scripts []string
	var downloadTarget = "/tmp/apulisedge"
	var packageName = "apulisedge"
	var fileName = packageName + ".tar.gz"
	var pubKeyFileName = packageName + ".public"
	var signFileName = packageName + ".sig"

	clu := configs.Config.Cluster
	fmt.Println("Host: ", configs.Config.Cluster.Host)
	fmt.Println("Domain: ", configs.Config.Cluster.Domain)

	// if configs.Config.Cluster.Host != "" {
	// 	address := net.ParseIP(configs.Config.Cluster.Host)
	// 	if address == nil {
	// 		// 没有匹配上
	// 		s, err := utils.GetIpByDns(configs.Config.Cluster.Host)
	// 		if err != nil || s == "" {
	// 			return nil, errors.New("域名解析错误")
	// 		}
	// 		ipCluster = s
	// 		fmt.Println(configs.Config.Cluster.Host)
	// 	} else {
	// 		// 匹配上
	// 		ipCluster = address.String()
	// 		fmt.Println(address)
	// 	}
	// } else {
	// 	return nil, errors.New("没有发现Host")
	// }

	step1 := ""
	if clu.Host != "" {
		address := net.ParseIP(clu.Host)
		if address == nil {
			// 域名
			step1 = "echo 'hello'"
		} else {
			// ip
			step1 = fmt.Sprintf("echo '%s %s' >> /etc/hosts", clu.Host, clu.Domain)
		}
	} else {
		return nil, errors.New("没有发现Host")
	}

	step1 = step1 + " && mkdir -p /tmp/apulisedge && tar -zxvf cert.tar.gz -C /tmp/apulisedge"
	step1 = step1 + " && mkdir -p /etc/kubeedge/certs"
	step1 = step1 + " && mkdir -p /etc/kubeedge/ca"
	step1 = step1 + " && mv /tmp/apulisedge/rootCA.crt /etc/kubeedge/ca"
	step1 = step1 + " && mv /tmp/apulisedge/server.key /etc/kubeedge/certs"
	step1 = step1 + " && mv /tmp/apulisedge/server.crt /etc/kubeedge/certs"
	scripts = append(scripts, step1)

	// step2
	// clean env
	step2 := "rm -rf " + downloadTarget + " && "
	// make neccessary dir
	step2 = step2 + " mkdir -p " + downloadTarget + " && "
	step2 = step2 + " mkdir -p /opt/apulisedge && "
	// download package and signature
	step2 = step2 + "wget " + configs.Config.Relevant.ApedgePublicKey + " -O /tmp/apulisedge/apedgenode.tar.gz && tar -zxvf /tmp/apulisedge/apedgenode.tar.gz -C /tmp/apulisedge && "
	// verify file
	step2 = step2 + " openssl dgst -verify " + downloadTarget + "/" + pubKeyFileName + " -sha256 -signature " + downloadTarget + "/" + signFileName + " " + downloadTarget + "/" + fileName + " && "
	// decompress package
	step2 = step2 + "tar -zxvf " + downloadTarget + "/" + fileName + " -C " + downloadTarget + " && "
	// move install script
	step2 = step2 + "cp -f " + downloadTarget + "/package/scripts/* /opt/apulisedge/" + " && "
	// run install script
	step2 = step2 + "bash /opt/apulisedge/install_edge.sh -d " + clu.Domain +
		" -l " + fmt.Sprintf("%s:%s", clu.EdgeImage, clu.EdgeImageTag) +
		" -h " + NodeId

	scripts = append(scripts, step2)
	scripts = append(scripts, "docker run -d --restart==always --privileged=true -v  /usr/local/Ascend/driver:/usr/local/Ascend/driver  -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi -v /var/lib/kubelet/device-plugins:/var/lib/kubelet/device-plugins  -v /var/log/devicePlugin:/var/log/devicePlugin harbor.apulis.cn:8443/release/ascend-k8sdeviceplugin-310:v1.1 ascendplugin --mode=ascend310 --useAscendDocker=false")
	scripts = append(scripts, "docker run -d --restart==always --privileged=true -v /var/lib/kubelet/device-plugins:/var/lib/kubelet/device-plugins -v /var/log/devicePlugin:/var/log/devicePlugin nvcr.io/nvidia/k8s-device-plugin nvidia-device-plugin --mig-strategy=none --pass-device-specs=false --fail-on-init-error=true --device-list-strategy=envvar --device-id-strategy=uuid --nvidia-driver-root=/")
	return scripts, nil
}

func CheckNodeIsExistByNodeID(tx *gorm.DB, nodeId string) (*dao.ApedgeNode, bool, error) {
	node := &dao.ApedgeNode{
		NodeId: nodeId,
	}
	_, err := node.GetNodeByNodeId(tx)
	if err != nil {
		tx.Rollback()
		if err.Error() == common.DB_ERR_RECORD_NOT_FOUND {
			return nil, false, nil
		}
		return nil, false, err
	}
	return node, true, nil
}

func CheckIsModelsEmpty(models []dto.Model) bool {
	return len(models) == 0
}

func CheckIsModelNumOver(models []dto.Model) bool {
	return len(models) > int(configs.Config.Relevant.ModelMaxNumPerTask)
}

func GenServiceId() string {
	timeN := time.Now().UnixNano()
	return fmt.Sprintf("edgeinfer-%s", strconv.FormatInt(timeN, 10))
}

// 获取模型信息
func ModelInfo(models []dto.Model, typeSpec string, serviceId string) ([]dao.ApflowModel, bool, error) {
	var apModels []dao.ApflowModel

	canStart, err := ModelCheckAndTransform(models, typeSpec, serviceId)
	if err != nil {
		return nil, canStart, err
	}
	for _, m := range models {
		modelPath, err := remote.GetModelPathFromApWorkShop(m.ID, m.VersionId)
		if err != nil {
			logger.Errorln("模型工厂请求模型地址错误：", err)
			return nil, canStart, common.ErrModelGetPath
		}
		apModels = append(apModels, dao.ApflowModel{
			Name:         m.Name,
			ID:           m.ID,
			Version:      m.Version,
			VersionId:    m.VersionId,
			PvcPath:      modelPath.StoragePath,
			DownloadLink: configs.Config.Cluster.Ip + modelPath.DownloadLink,
		})
	}
	return apModels, canStart, nil
}

// 模型检查和转换 models
func ModelCheckAndTransform(models []dto.Model, typeSpec string, serviceId string) (bool, error) {

	for _, m := range models {
		// 检查
		modelCheck, err := remote.ModelCheckFromApWorkShop(m.ID, m.VersionId, typeSpec)
		if err != nil {
			logger.Errorln(err)
			return false, common.ErrModelCheck
		}
		// 不能使用
		if !modelCheck.CanInference {
			return false, common.ErrModelCanNotInfer
		}
		// 不需要转换
		if !modelCheck.NeedTransform {
			return true, nil
		} else {
			// 需要转换
			_, err = remote.ModelTransformApWorkShop(m.ID, m.VersionId, typeSpec, serviceId, configs.Config.Rabbitmq.QueueModelTransform)
			if err != nil {
				logger.Errorln(err)
				return false, common.ErrModelTransform
			}
		}
	}
	return false, nil
}

func StartApflow(info dao.ApedgeInference, snInfo dao.ApedgeServiceNode) error {
	logger.Infoln("使用镜像 ", fmt.Sprintf("%s:%s", info.ContainerImage, info.ContainerImageTag))
	if info.ContainerImage == configs.Config.Relevant.Image.TritonImage {
		logger.Infoln("创建GPU推理")
		infer := &remote.GpuApedgeInference{}
		_, err := infer.Creat(info, snInfo)
		if err != nil {
			return err
		}
	} else if info.ContainerImage == configs.Config.Relevant.Image.NpuInferImage {
		logger.Infoln("创建NPU推理")
		infer := &remote.NpuApedgeInference{}
		_, err := infer.Creat(info, snInfo)
		if err != nil {
			return err
		}
	} else {
		return errors.New("镜像不能用于推理，请重新选择")
	}

	return nil
}

func CheckCanStopService(status string) bool {
	if status == common.TASK_STATUS_RUNNING || status == common.TASK_STATUS_SCHEDULING {
		return true
	}
	return false
}

func CheckCanStartService(status string) bool {
	if status == common.TASK_STATUS_MODEL_TRANSFORM_OK || status == common.TASK_STATUS_STOP {
		return true
	}
	return false
}
